from collections import OrderedDict

import torch
import torch.nn.functional as F
from torch import nn

class LayerNorm(nn.LayerNorm):
    """Subclass torch's LayerNorm to handle fp16."""

    def forward(self, x: torch.Tensor):
        orig_type = x.dtype
        ret = super().forward(x.type(torch.float32))
        return ret.type(orig_type)


class QuickGELU(nn.Module):
    def forward(self, x: torch.Tensor):
        return x * torch.sigmoid(1.702 * x)



class ResidualAttentionBlock(nn.Module):
    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
        super().__init__()

        self.attn = nn.MultiheadAttention(d_model, n_head)
        self.ln_1 = LayerNorm(d_model)
        self.mlp = nn.Sequential(OrderedDict([
            ("c_fc", nn.Linear(d_model, d_model * 4)),
            ("gelu", QuickGELU()),
            ("c_proj", nn.Linear(d_model * 4, d_model))
        ]))
        self.ln_2 = LayerNorm(d_model)
        self.attn_mask = attn_mask

    def attention(self, x: torch.Tensor):
        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
        return self.attn(x, x, x, need_weights=True, attn_mask=self.attn_mask)

    def forward(self, x):
        x = x[0]
        tgt, attn_map = self.attention(self.ln_1(x))
        x = x + tgt
        x = x + self.mlp(self.ln_2(x))
        return [x, attn_map]


class Transformer(nn.Module):
    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None, drop_path_rate=0.):
        super().__init__()
        self.width = width
        self.layers = layers
        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])

    def forward(self, x: torch.Tensor):
        x, attn_map = self.resblocks([x])[0], self.resblocks([x])[1]
        return x, attn_map


class VisionTransformer(nn.Module):
    def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int):
        super().__init__()
        self.input_resolution = input_resolution
        self.output_dim = output_dim
        #   容易inf参数化
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)

        scale = width ** -0.5
        self.class_embedding = nn.Parameter(scale * torch.randn(width))
        self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
        self.ln_pre = LayerNorm(width)
        #   容易inf参数化
        self.transformer = Transformer(width, layers, heads)

        self.ln_post = LayerNorm(width)
         #   容易inf参数化
        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))

    def forward(self, x: torch.Tensor):
        # B, C, H, W = x.shape
        #   1, 3, 224, 224
        x = self.conv1(x)  # shape = [*, width, grid, grid]
        B, C, H, W = x.shape
        #   1, 768, 7, 7
        x = x.reshape(B, C, H*W)  # shape = [*, width, grid ** 2]
        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
        #   1, 50, 768
        zero_x = torch.zeros(B, 1, C, dtype=x.dtype, device=x.device)
        class_embedding = zero_x + self.class_embedding.to(x.dtype)
        x = torch.cat([class_embedding,  x], dim=1)  # shape = [*, grid ** 2 + 1, width]
        
        
        positional_embedding = self.positional_embedding.to(x.dtype)
        x = x + positional_embedding
        x = self.ln_pre(x)
        x = x.permute(1, 0, 2)  # NLD -> LND
        
        x, attn_map = self.transformer(x)
        mem_x = x.permute(1, 0, 2)  # LND -> NLD
        #   1, 50, 768
        
        #x = self.ln_post(mem_x[:, :1, :])#这里进行更改了
        
        x = self.ln_post(mem_x)
        #   1, 768
        if self.proj is not None:
            x = x @ self.proj
            
        x = x.permute(1, 0, 2)
        #   197, 8, 512
        return x, attn_map, mem_x


